package org.java;

import com.alibaba.excel.util.StringUtils;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.apache.sshd.client.SshClient;
import org.apache.sshd.client.channel.ChannelExec;
import org.apache.sshd.client.channel.ClientChannelEvent;
import org.apache.sshd.client.future.ConnectFuture;
import org.apache.sshd.client.keyverifier.DefaultKnownHostsServerKeyVerifier;
import org.apache.sshd.client.keyverifier.KnownHostsServerKeyVerifier;
import org.apache.sshd.client.keyverifier.ServerKeyVerifier;
import org.apache.sshd.client.session.ClientSession;
import org.apache.sshd.common.util.security.SecurityUtils;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.ObjectOutputStream;
import java.net.SocketAddress;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.util.EnumSet;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * @Description Java特性测试类
 * @Date 2024/8/13 10:11
 * @Created by 10574
 */
public class RemoteExecuteUtil {
    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    public static class Device {
        private double freeMemory;

        private double totalMemory;

        private double usedMemory;

        private int temperature;

        private int powerStatus;
    }

    private static final String key = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAICsYcbaupMa4Ak31IiPWgZKnksQd94s4CGF7+n7JkZrA root@master";
    public static void main(String[] args) throws Exception {
        Map<String, Device> devices = new HashMap<>();
        final double nodeTotalGpuUsage = devices.values().stream().mapToDouble(Device::getTotalMemory).sum();
        System.out.println(nodeTotalGpuUsage);

    }

    private String ip;

    private Integer port;

    private String username;

    private String password;


    public RemoteExecuteUtil(String ip, Integer port, String username, String password) {
        this.ip = ip;
        this.port = port;
        this.username = username;
        this.password = password;
    }

    /**
     * 连接远程主机
     *
     * @return
     * @throws IOException
     */
    private Boolean login(SshClient client) throws Exception {
        client.start();
        KnownHostsServerKeyVerifier verifier = new DefaultKnownHostsServerKeyVerifier(new ServerKeyVerifier() {
            @Override
            public boolean verifyServerKey(ClientSession clientSession, SocketAddress remoteAddress,
                                           PublicKey serverKey) {
                return true;
            }
        }, true);
        client.setServerKeyVerifier(verifier);
        client.start();
        return client.isStarted();
    }
    /**
     * 获取密钥对
     */
    static KeyPair getKeyPairFromString(String pk) throws IOException, NoSuchAlgorithmException {
        final KeyPairGenerator rsa = KeyPairGenerator.getInstance("ed25519");
        final KeyPair keyPair = rsa.generateKeyPair();
        final ByteArrayOutputStream stream = new ByteArrayOutputStream();
        stream.write(pk.getBytes());
        final ObjectOutputStream o = new ObjectOutputStream(stream);
        o.writeObject(keyPair);
        return keyPair;
    }
    /**
     * 远程执行ssh命令
     *
     * @param command
     * @return
     * @throws IOException
     */
    public String execute(String command) throws Exception {
        String result = "";
        String errResult = "";
        SshClient client = SshClient.setUpDefaultClient();
        try {
            if (login(client)) {
                ConnectFuture cf = client.connect(username, ip, port);
                ClientSession session = cf.verify().getSession();
                session.addPublicKeyIdentity(getKeyPairFromString(key));
                session.auth().verify(10000L);
//                session.addPasswordIdentity(password);
                session.auth().verify(TimeUnit.SECONDS.toMillis(30));
//                System.out.println("connect ssh success");
//                System.out.println("cmd >>> " + command);
                ChannelExec ce = session.createExecChannel(command);
                ByteArrayOutputStream out = new ByteArrayOutputStream();
                ByteArrayOutputStream err = new ByteArrayOutputStream();
                ce.setOut(out);
                ce.setErr(err);
                ce.open();
                ce.waitFor(EnumSet.of(ClientChannelEvent.CLOSED), -1);

                result = out.toString();
                errResult = err.toString();
                out.close();
                err.close();
//                System.out.println("result >>>>> " + result);
                if (StringUtils.isBlank(result)) {
                    if (!StringUtils.isBlank(errResult)) {
//                        System.out.println("err>>>>>>>>>" + errResult);
                        throw new RuntimeException(errResult);
                    }
                }
            }
        } finally {
            if (client != null) {
                client.stop();
            }
        }
        return result;
    }

}